import cv2
import numpy as np
from scipy import ndimage

# 读取两张图片
image1 = cv2.imread('out_img/Img_0.png')
image2 = cv2.imread('template_640.jpg')

# 确保两张图片大小一致
assert image1.shape == image2.shape, "两张图片的尺寸不一致"

# 转换图片为RGB格式
image1_rgb = cv2.cvtColor(image1, cv2.COLOR_BGR2RGB)
image2_rgb = cv2.cvtColor(image2, cv2.COLOR_BGR2RGB)

# 定义一个白色区域的阈值 (假设白色为RGB(255, 255, 255))
white_threshold = [240, 240, 240]

# 获取第一张图片中白色区域的掩模
mask_white_area = np.all(image1_rgb >= white_threshold, axis=-1)

# 获取第二张图片在相应位置的区域
image2_white_area = np.copy(image2_rgb)
image2_white_area[~mask_white_area] = 0  # 保留白色区域，其他区域变为黑色

# 对白色区域进行泊松融合
def poisson_blend(source_img, target_img, mask):
    # 对每个颜色通道分别进行拉普拉斯操作
    blended_image = np.zeros_like(source_img)

    for i in range(3):  # RGB通道
        # 计算单通道的拉普拉斯算子
        laplacian = ndimage.laplace(source_img[:, :, i])

        # 进行融合
        blended_image[:, :, i] = source_img[:, :, i] + laplacian * mask

    return np.clip(blended_image, 0, 255).astype(np.uint8)

# 使用泊松分布融合图像
blended_white_area = poisson_blend(image2_white_area, image1_rgb, mask_white_area)

# 将第一张图的非白色部分直接贴到第二张图的相应位置
blended_image = np.copy(image2_rgb)
blended_image[mask_white_area] = blended_white_area[mask_white_area]  # 用融合后的白色区域替换
blended_image[~mask_white_area] = image1_rgb[~mask_white_area]  # 将非白色区域从第一张图复制过来

# 保存结果
blended_image_bgr = cv2.cvtColor(blended_image, cv2.COLOR_RGB2BGR)
cv2.imwrite('blended_image.jpg', blended_image_bgr)

# 显示结果
cv2.imshow('Blended Image', blended_image_bgr)
cv2.waitKey(0)
cv2.destroyAllWindows()